'''
Created on Nov 27, 2010

@author: oabalbin
'''
import subprocess
import exome.gatk_cluster.samtools_commands_cluster  as mysam
import exome.gatk_cluster.picard_commands_cluster as mypicard
import exome.gatk_cluster.cluster_jobs_header as jh
from collections import defaultdict, deque

def reads_alignment(ref_genome, reads_file_name, nmismatch, num_cores, outfile_name, path_to_bwa):
    '''
    time /exds/users/oabalbin/sw/bwa/bwa-0.5.8a/bwa aln -k 2 -t 3 
    /exds/projects/alignment_indexes/bwa/hg19_illumina/hg19_illumina.fa s_4_2_sequence.txt > 
    s_4_2_sequence.txt.sai
    '''
    #outfile_name = reads_file_name.replace+'.sai'
    bwa_command = path_to_bwa+'bwa'
    args=[bwa_command,'aln', '-k', str(nmismatch), '-t', str(num_cores), ref_genome, reads_file_name, '>', outfile_name]
        
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command


def paired_read_alignment_sam(ref_genome, align_reads_pair1, align_reads_pair2, 
                              raw_reads1, raw_reads2, outfile_name, path_to_bwa):
    '''
    time /exds/users/oabalbin/sw/bwa/bwa-0.5.8a/bwa sampe 
    /exds/projects/alignment_indexes/bwa/hg19_illumina/hg19_illumina.fa 
    ./s_4_1_sequence.txt.sai ./s_4_2_sequence.txt.sai ./s_4_1_sequence.txt 
    ./s_4_2_sequence.txt > s_4_12_sequence.aln.sam
    
    # Using the path for the gatk well formed bam files with read group, 
    and platform information
    time /exds/users/oabalbin/sw/bwa/bwa-0.5.8a/bwa sampe 
    -i MARK -m DEPRISTO -l LIBRARY -p ILLUMINA 
    /exds/projects/alignment_indexes/gatk/hg19/bwa/hg19.fa 
    ./s_3_1_sequence.hg19.sai ./s_3_2_sequence.hg19.sai 
    ./s_3_1_sequence.txt ./s_3_2_sequence.txt > s_3_12_sequence.hg19.aln2.sam
    
    '''
    #fields=align_reads_pair1.split('_')
    #name = fields[0]+'_'+fields[1]+'_'+'12'+fields[3]
    #outfile_name=name.replace('.sai','.sam')
    bwa_command=path_to_bwa+'bwa'
    args=[bwa_command, 'sampe', '-i', 'mctp', '-m', 'sptest', '-l', 'LIBRARY', '-p', 'ILLUMINA', 
          ref_genome, align_reads_pair1,align_reads_pair2, raw_reads1, raw_reads2, '>', outfile_name]
    
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command



def main_bwa_alignment(sample_dict, bwa_run_dict):
    '''
    It perfomrs a paired bwa alignment
    remove duplicate, sort the reads in cordinate order, and index the sorted file
    bwa_param => dictionary with the parameters for a bea run
    nmismath= number of mismatches
    '''
    ref_genome=bwa_run_dict['ref_genome_bwa']
    nmismatch, num_cores= bwa_run_dict['nmismatches'], bwa_run_dict['num_cores'] 
    path_to_bwa, path_to_sam=bwa_run_dict['path_to_bwa'],bwa_run_dict['path_to_sam']
    path_to_picard=bwa_run_dict['path_to_picard']
    my_email=bwa_run_dict['myemail']
    
    
    node_memory = 45000.0
    node_processors = 12
    single_processor=1
    mem_per_processors = int(float(node_memory) / node_processors)
    extra_mem=8000
    # wt=walltime
    short_wtime = "24:00:00"
    long_wtime="100:00:00"

    wt_aligning="60:00:00"
    wt_samtools="24:00:00"

    
    
    new_sample_dict=defaultdict(list)
    
    for sample_name, read_pairs_set in sample_dict.iteritems():
        aligned_reads=deque()
        aln_jobs=deque()
        for thispair in read_pairs_set:
            print sample_name, thispair
            jname=sample_name
            outfile_name = thispair + '.sai'
            command = reads_alignment(ref_genome, thispair, nmismatch, num_cores, outfile_name, path_to_bwa)
            
            jobidaln = jh.qsub(jname+'_anl', command, num_cores, cwd=None, walltime=wt_aligning, pmem=None, 
                                  deps=None, stdout=None, email_addresses=my_email)
            
            aligned_reads.append(outfile_name)
            aln_jobs.append(jobidaln)
            
        #sampe = paired aligned file (sam), bampe paired aligned file (bam)
        # brs = bampe removed duplicates(r) and coordinate sorted (s) file
        fields=aligned_reads[0].split('_')
        name = fields[0]+'_'+fields[1]+'_'+'12_'+fields[3]
        sampe_file_name=name.replace('.sai','.sam')

        command = paired_read_alignment_sam(ref_genome, aligned_reads[0],aligned_reads[1], 
                                  read_pairs_set[0], read_pairs_set[1], sampe_file_name, path_to_bwa)
        
        jobidpe = jh.qsub(jname+'_pe', command, single_processor, cwd=None, walltime=wt_aligning, pmem=extra_mem, 
                                  deps=list(aln_jobs), stdout=None, email_addresses=my_email)
        
        
        command, bampe_file = mysam.sam2bam(ref_genome, sampe_file_name,path_to_sam)
        
        jobidsb = jh.qsub(jname+'_sb', command, single_processor, cwd=None, walltime=wt_samtools, pmem=extra_mem, 
                                  deps=jobidpe, stdout=None, email_addresses=my_email)

        # remove duplicates. Duplicates are not going
        # to be removed but marked with picard later on
        #bampe_rmdup = mysam.remove_duplicates(bampe_file)
        
        # Sort bam file
        #command, bampe_rmdup_sorted = mysam.sort_bam_file(bampe_file,path_to_sam)
        
        command, bampe_rmdup_sorted = mypicard.sortSam(bampe_file,extra_mem, path_to_picard)
        jobidbs = jh.qsub(jname+'_bs', command, single_processor, cwd=None, walltime=wt_samtools, pmem=extra_mem, 
                                  deps=jobidsb, stdout=None, email_addresses=my_email)
        
        # Index bam file
        command, brs_indexed = mysam.index_bam_file(bampe_rmdup_sorted, path_to_sam)
        jobidbi = jh.qsub(jname+'_bi', command, single_processor, cwd=None, walltime=wt_samtools, pmem=extra_mem, 
                                  deps=jobidbs, stdout=None, email_addresses=my_email)
        
        new_sample_dict[sample_name].append(bampe_rmdup_sorted)

    
    # this function should return a dictionary of sample_lanes key:sample, value: indexed bam file of pair end alignment
    return jobidbi, new_sample_dict


if __name__ == '__main__':
         
    sample_lanes_dict={'aM18':['/nobackup/med-mctp/oabalbin/test/s_3_1_sequence.txt','/nobackup/med-mctp/oabalbin/test/s_3_2_sequence.txt'],
                       'aM18_N':['/nobackup/med-mctp/oabalbin/test/s_4_1_sequence.txt','/nobackup/med-mctp/oabalbin/test/s_4_2_sequence.txt']}
    
    run_dict = {'path_to_gatk':'/nobackup/med-mctp/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/',
                     'path_to_picard':'/nobackup/med-mctp/sw/bioinfo/picard/picard-tools-1.35/', 
                     'path_to_bwa':'/nobackup/med-mctp/sw/bioinfo/alignment/bwa/bwa-0.5.8a/',
                     'path_to_sam':'/nobackup/med-mctp/sw/bioinfo/samtools/samtools-0.1.10/',
                     'resources_folder':'/nobackup/med-mctp/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources/', 
                     'rscipt_path':'/home/software/rhel5/R/2.10.1-gcc/bin/Rscript',
                     'use_mem':45000.0, 'num_cores':6,
                     'nmismatches':2,
                     'ref_genome':'/nobackup/med-mctp/sw/alignment_indexes/gatk/hg19/hg19.fa',
                     'ref_genome_bwa':'/nobackup/med-mctp/sw/alignment_indexes/bwa/hg19/hg19.fa',
                     'snpdb_file':'/nobackup/med-mctp/sw/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed.vcf',
                     'indeldb_file':'/nobackup/med-mctp/sw/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed.vcf',
                     'path_to_intervals':'/nobackup/med-mctp/oabalbin/test/',
                     'recal_analysis_outputdir':'/nobackup/med-mctp/oabalbin/test/recal_analysis/',
                     'temp_dir':'/nobackup/med-mctp/oabalbin/test/temp/',
                     'qsubfile':'/nobackup/med-mctp/oabalbin/test/',
                     'out_dir':'/nobackup/med-mctp/oabalbin/test/',
                     'myemail':['alebalbin@gmail.com']
                     }

    jobidbi, new_sample_dict =main_bwa_alignment(sample_lanes_dict, run_dict)
    
    
    
    